# 获取初始设定的pose, velocity, tau
from os.path import dirname, join, abspath
import yaml
import numpy as np
import pinocchio

# 这里open的基路径是main.py
with open(join(dirname(str(abspath(__file__))),'init.yml'), 'r', encoding='utf-8') as f:
    init_config= yaml.load(f.read(), Loader=yaml.FullLoader)

def pvtInit(model, ROBOT = False):

    # 根据Robot关键字的bool值，选取特定的init.yml文件
    if ROBOT:
        # 这里open的基路径是main.py
        with open(join(dirname(str(abspath(__file__))),'robot_init.yml'), 'r', encoding='utf-8') as f:
            init_config= yaml.load(f.read(), Loader=yaml.FullLoader)
    else:
        with open(join(dirname(str(abspath(__file__))),'init.yml'), 'r', encoding='utf-8') as f:
            init_config= yaml.load(f.read(), Loader=yaml.FullLoader)
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    
    if init_config['zero_p']:
        p0 = pinocchio.neutral(model)
        
    elif init_config['random_p']:
        np.random.seed(init_config['seed'])
        p0 = np.random.randn(model.nq)
    else:
        p0_list = init_config.get('p0',[])
        p0 = np.array(p0_list, dtype=float)

    if init_config['zero_v']:
        v0 = np.zeros(model.nv)
    elif init_config['random_v']:
        np.random.seed(init_config['seed'])
        v0 = np.random.randn(model.nv)*0.1
    else:
        v0 = np.array(init_config.get('v0',[]))

    if init_config['zero_tau']:
        tau0 = np.zeros(model.nv)
    elif init_config['random_tau']:
        np.random.seed(init_config['seed'])
        tau0 = np.random.randn(model.nv)*0.01
    else:
        tau0 = np.array(init_config.get('tau0',[]))

    print('tau0 is ', tau0)

    return p0, v0, tau0
